/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#ifndef __REGRESSION_WORKER_HH_
#define __REGRESSION_WORKER_HH_

#include <boost/variant.hpp>
#include <boost/thread.hpp>

#include "example.hh"
#include "regressor.hh"

namespace flassol
{
  class Worker
    {
      private:
        typedef Regressor::UpdateInfo UpdateInfo;

        struct EstimateOperation
          {
            const Regressor*                            regressor;
            const IndexedFeatureSet*                    f;
            uint64_t                                    thread_no;
            uint64_t                                    num_threads;

            volatile bool                               complete;
            std::vector<Regressor::EstimateResult>      result;

            EstimateOperation (const Regressor&         _regressor,
                               const IndexedFeatureSet* _f,
                               uint64_t                 _thread_no,
                               uint64_t                 _num_threads,
                               uint64_t                 _num_offset)
              : regressor (&_regressor),
                f (_f),
                thread_no (_thread_no),
                num_threads (_num_threads),
                complete (false),
                result (_num_offset)
              {
              }
          };

        struct AdaptiveNormOperation
          {
            const Regressor*              regressor;
            const IndexedFeatureSet*      f;
            uint64_t                      thread_no;
            uint64_t                      num_threads;

            volatile bool                 complete;
            std::vector<UpdateInfo>       result;

            AdaptiveNormOperation (const Regressor&         _regressor,
                                   const IndexedFeatureSet* _f,
                                   uint64_t                 _thread_no,
                                   uint64_t                 _num_threads,
                                   uint64_t                 _num_offset)
              : regressor (&_regressor),
                f (_f),
                thread_no (_thread_no),
                num_threads (_num_threads),
                complete (false),
                result (_num_offset)
              {
              }
          };

        struct UpdateOperation
          {
            Regressor*                          regressor;
            const IndexedFeatureSet*            f;
            const std::vector<UpdateInfo>*      info;
            uint64_t                            thread_no;
            uint64_t                            num_threads;

            volatile bool                       complete;

            UpdateOperation (Regressor&                     _regressor,
                             const IndexedFeatureSet*       _f,
                             const std::vector<UpdateInfo>& _info,
                             uint64_t                       _thread_no,
                             uint64_t                       _num_threads)
              : regressor (&_regressor),
                f (_f),
                info (&_info),
                thread_no (_thread_no),
                num_threads (_num_threads),
                complete (false)
              {
              }
          };

        struct NullOperation
          {
            volatile bool           complete;
          };

        typedef boost::variant<EstimateOperation, 
                               UpdateOperation,
                               AdaptiveNormOperation,
                               NullOperation> Operation;

        volatile bool data_ready;
        boost::condition_variable send;
        boost::condition_variable recv;
        boost::mutex mutex;
        boost::thread thread;
        Operation operation;

        class Loop
          {
            private:
              Worker* worker;

              class OperationVisitor : public boost::static_visitor<>
                {
                  private:
                    Worker* worker;

                  public:
                    explicit 
                    OperationVisitor (Worker* _worker) : worker (_worker) { }

                    void
                    operator() (EstimateOperation& op) const
                      {
                        op.regressor->get_partial (op.f,
                                                   op.thread_no,
                                                   op.num_threads,
                                                   op.result);
      
                        op.complete = true;
                      }

                    void
                    operator() (AdaptiveNormOperation& op) const
                      {
                        op.regressor->get_adaptive_norm_partial (op.f,
                                                                 op.thread_no,
                                                                 op.num_threads,
                                                                 op.result);

                        op.complete = true;
                      }

                    void
                    operator() (UpdateOperation& op) const
                      {
                        op.regressor->update_partial (op.f,
                                                      *op.info,
                                                      op.thread_no,
                                                      op.num_threads);
      
                        op.complete = true;
                      }

                    void
                    operator() (NullOperation& op) const
                      {
                        op.complete = true;
                      }
                };

            public:
              explicit Loop (Worker* _worker) : worker (_worker) { } 

              void 
              operator () ()
                {
                  while (1)
                    {
                      boost::unique_lock<boost::mutex> lock (worker->mutex);

                      while (! worker->data_ready)
                        {
                          worker->send.wait (lock);
                        }

                      worker->data_ready = false;

                      boost::apply_visitor (OperationVisitor (worker),
                                            worker->operation);

                      lock.unlock ();

                      worker->recv.notify_one ();
                    }
                }
          };

      public:
        explicit
        Worker () : data_ready (false),
                    send (),
                    recv (),
                    mutex (),
                    thread (boost::thread (Loop (this))),
                    operation (NullOperation ())
          {
          }

        void
        estimate_start (const Regressor&         r,
                        const IndexedFeatureSet* f,
                        uint64_t                 thread_no,
                        uint64_t                 num_threads,
                        uint64_t                 num_offset)
          {
            EstimateOperation op (r, f, thread_no, num_threads, num_offset);

            {
              boost::lock_guard<boost::mutex> lock (mutex);
              operation = op;
              data_ready = true;
            }

            send.notify_one ();
          }

        void
        adaptive_norm_start (const Regressor&               r,
                             const IndexedFeatureSet*       f,
                             uint64_t                       thread_no,
                             uint64_t                       num_threads,
                             const std::vector<UpdateInfo>& info)
          {
            AdaptiveNormOperation op (r, f, thread_no, num_threads, info.size ());

            for (unsigned int n = 0; n < info.size (); ++n)
              {
                op.result[n].gsq = info[n].gsq;
              }

            {
              boost::lock_guard<boost::mutex> lock (mutex);
              operation = op;
              data_ready = true;
            }

            send.notify_one ();
          }

        void
        update_start (Regressor&                        r,
                      const IndexedFeatureSet*          f,
                      const std::vector<UpdateInfo>&    info,
                      uint64_t                          thread_no,
                      uint64_t                          num_threads)
          {
            UpdateOperation op (r, f, info, thread_no, num_threads);
            {
              boost::lock_guard<boost::mutex> lock (mutex);
              operation = op;
              data_ready = true;
            }
            send.notify_one ();
          }

        void
        estimate_finish (std::vector<Regressor::EstimateResult>& rv)
          {
            boost::unique_lock<boost::mutex> lock (mutex);

            while (! boost::get<EstimateOperation> (operation).complete)
              {
                recv.wait (lock);
              }

            for (uint64_t n = 0; n < rv.size (); ++n)
              {
                rv[n] += boost::get<EstimateOperation> (operation).result[n];
              }

            lock.unlock ();
          }

        void
        adaptive_norm_finish (std::vector<Regressor::UpdateInfo>& rv)
          {
            boost::unique_lock<boost::mutex> lock (mutex);

            while (! boost::get<AdaptiveNormOperation> (operation).complete)
              {
                recv.wait (lock);
              }

            for (uint64_t n = 0; n < rv.size (); ++n)
              {
                rv[n].adaptivexnorm += 
                  boost::get<AdaptiveNormOperation> (operation)
                    .result[n].adaptivexnorm;
              }

            lock.unlock ();
          }

        void
        update_finish ()
          {
            boost::unique_lock<boost::mutex> lock (mutex);

            while (! boost::get<UpdateOperation> (operation).complete)
              {
                recv.wait (lock);
              }

            lock.unlock ();
          }

        ~Worker ()
          {
            thread.interrupt ();
            thread.join ();
          }
    };

  class WorkerPool 
    {
      private:
        typedef Regressor::UpdateInfo UpdateInfo;

        unsigned int num_threads;
        Worker* workers; // std::vector and boost::thread 
                         // without rvalue references ... 
                         // how can it work ?
                         //

      public:
        explicit
        WorkerPool (unsigned int num_extra_threads) 
          : num_threads (1 + num_extra_threads),
            workers (new Worker[num_threads])
          {
          }

        void
        estimate (const Regressor&                        r,
                  const IndexedFeatureSet*                f,
                  std::vector<Regressor::EstimateResult>& rv)
          {
            for (uint64_t t = 1; t < num_threads; ++t)
              {
                workers[t].estimate_start (r, f, t, num_threads, rv.size ());
              }

            r.get_partial (f, 0, num_threads, rv);

            for (uint64_t t = 1; t < num_threads; ++t)
              {
                workers[t].estimate_finish (rv);
              }
          }

        void
        adaptive_norm (const Regressor&                    r,
                       const IndexedFeatureSet*            f,
                       std::vector<Regressor::UpdateInfo>& rv)
          {
            for (uint64_t t = 1; t < num_threads; ++t)
              {
                workers[t].adaptive_norm_start (r, f, t, num_threads, rv);
              }

            r.get_adaptive_norm_partial (f, 0, num_threads, rv);

            for (uint64_t t = 1; t < num_threads; ++t)
              {
                workers[t].adaptive_norm_finish (rv);
              }
          }

        void
        update (Regressor&                      r,
                const IndexedFeatureSet*        f,
                const std::vector<UpdateInfo>&  info)
          {
            for (uint64_t t = 1; t < num_threads; ++t)
              {
                workers[t].update_start (r, f, info, t, num_threads);
              }

            r.update_partial (f, info, 0, num_threads);

            for (uint64_t t = 1; t < num_threads; ++t)
              {
                workers[t].update_finish ();
              }
          }

        ~WorkerPool ()
          {
            delete[] workers;
          }
    };
}

#endif // __REGRESSION_WORKER_HH_

